extern "C" {
#include <immintrin.h>
}

#define A(i, j) A[(i) + (j) * (LDA)]
#define B(i, j) B[(i) + (j) * (LDB)]
#define C(i, j) C[(i) + (j) * (LDC)]
#define KERNEL_K1                                                              \
    do {                                                                       \
        __m256d a = _mm256_loadu_pd(&A(i, k));                                 \
        __m256d b0 = _mm256_broadcast_sd(&B(k, j));                            \
        __m256d b1 = _mm256_broadcast_sd(&B(k, j + 1));                        \
        __m256d b2 = _mm256_broadcast_sd(&B(k, j + 2));                        \
        __m256d b3 = _mm256_broadcast_sd(&B(k, j + 3));                        \
        c0 = _mm256_fmadd_pd(a, b0, c0);                                       \
        c1 = _mm256_fmadd_pd(a, b1, c1);                                       \
        c2 = _mm256_fmadd_pd(a, b2, c2);                                       \
        c3 = _mm256_fmadd_pd(a, b3, c3);                                       \
    } while (0);

static void mm_v6_corner(int M, int N, int K, double *A, int LDA, double *B,
                         int LDB, double *C, int LDC) {
    double tmp = 0;
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            tmp = 0;
            for (int k = 0; k < K; k++) {
                tmp += A(i, k) * B(k, j);
            }
            C(i, j) = tmp;
        }
    }
}
// 使用 AVX 指令集 4x4分块
// A:[M, K] B:[K, N] C:[M, N]
inline void mm_v6(int M, int N, int K, double *A, int LDA, double *B, int LDB,
                  double *C, int LDC) {
    int m4 = M & -4;
    int n4 = N & -4;
    int k4 = K & -4;
    for (int i = 0; i < m4; i += 4) {
        for (int j = 0; j < n4; j += 4) {
            __m256d c0 = _mm256_setzero_pd();
            __m256d c1 = _mm256_setzero_pd();
            __m256d c2 = _mm256_setzero_pd();
            __m256d c3 = _mm256_setzero_pd();
            for (int k = 0; k < k4;) {
                KERNEL_K1;
                k++;
                KERNEL_K1;
                k++;
                KERNEL_K1;
                k++;
                KERNEL_K1;
                k++;
            }
            for (int k = k4; k < K;) {
                KERNEL_K1;
                k++;
            }
            _mm256_storeu_pd(&C(i, j), c0);
            _mm256_storeu_pd(&C(i, j + 1), c1);
            _mm256_storeu_pd(&C(i, j + 2), c2);
            _mm256_storeu_pd(&C(i, j + 3), c3);
        }
    }
    if (M == m4 && N == n4) {
        return;
    }
    if (M != m4) {
        mm_v6_corner(M - m4, N, K, &A(m4, 0), LDA, &B(0, 0), LDB, &C(m4, 0),
                     LDC);
    }
    if (N != n4) {
        mm_v6_corner(m4, N - n4, K, &A(0, 0), LDA, &B(0, n4), LDB, &C(0, n4),
                     LDC);
    }
}
#undef A
#undef B
#undef C
#undef KERNEL_K1
